# Author: Liwei Wang

"""        
Test case for dlsodis(). Demo from opkddemos.f  in ODEPACK
The Burgers equation.

    This program solves a semi-discretized form of the following system
    Of three PDEs (each similar to a Burgers equation):
        This program solves a semi-discretized form of the Burgers equation,
        u  = -(u*u/2)  + eta * u
         t           x          xx
        for  -1 .le. x .le. 1, t .ge. 0.
    Here eta = 0.05.
    Boundary conditions: u(-1,t) = u(1,t) and du/dx(-1,t) = du/dx(1,t).
    Initial profile: square wave
      u(0,x) = 0    for 1/2 .lt. abs(x) .le. 1
      u(0,x) = 1/2  for abs(x) = 1/2
      u(0,x) = 1    for 0 .le. abs(x) .lt. 1/2
    An ODE system is generated by a simplified Galerkin treatment of the spatial
    variable x.
"""
from odespy import *
import numpy as np

def res(y, t, s, ires):
    r = np.zeros(12, float)
    r[0] = 1.5*(y[11]**2 - y[1]**2) + 1.8*(y[1] - 2.*y[0] + y[11])
    r[1:11] = 1.5*(y[:10]**2 - y[2:12]**2) + 1.8*(y[2:12] - 2.*y[1:11] + y[:10])
    r[11] = 1.5*(y[11]**2 - y[0]**2) + 1.8*(y[0] - 2.*y[11] + y[10])
    if ires!=-1:
        r[0] -= (s[11] + 4.*s[0] + s[1])/6.
        r[1:11] -= (s[:10] + 4.*s[1:11] + s[2:12])/6.
        r[11] -= (s[10] + 4.*s[11] + s[0])/6.
    return r,ires

def adda(y, t, j, ia, ja, p):
    p[j] += 2./3.
    p[(j+1) % 12] += 1./6.
    p[(j-1) % 12] += 1./6.
    return p 

def jac(y, t, s, j, ia, ja):
    p = np.zeros(12, float)
    p[(j-1) % 12] = -3.*y[j] + 1.8
    p[j] = -3.6
    p[(j+1) % 12] = 3.*y[j] + 1.8
    return p

t0, tn, n_points = 0., .5, 6
u0 = [0.]*3 + [.5] + [1.]*5 + [.5] + [0.]*2
time_points = np.linspace(t0, tn, n_points)

atol, rtol = 1e-2, 1e-2
ia = [1,4,7,10,13,16,19,22,25,28,31,34,37]
ja = [12,1,2,1,2,3,2,3,4,3,4,5,4,5,6,5,6,7,6,7,8,7,8,9,8,
      9,10,9,10,11,10,11,12,11,12,1]
ic = [1,4,7,10,13,16,19,22,25,28,31,34,37]
jc = [12,1,2,1,2,3,2,3,4,3,4,5,4,5,6,5,6,7,6,7,8,7,8,9,8,
      9,10,9,10,11,10,11,12,11,12,1]

exact_final = [1.60581860e-02, 3.23063251e-02, 1.21903380e-01, 
               2.70943828e-01, 4.60951522e-01, 6.57571216e-01, 
               8.25154453e-01, 9.35644796e-01, 9.90167557e-01, 
               9.22421221e-01, 5.85764902e-01, 1.81112615e-01]

method = Lsodis
# Test case 1: with res, adda, ia, ja, ic, jc & jac
m = method(res=res, adda_lsodis=adda, atol=atol, rtol=rtol, jac_lsodis=jac,
           ia=ia, ja=ja, ic=ic, jc=jc)
m.set_initial_condition(u0)
y, t = m.solve(time_points)
print 'Max error for test case 1 is %g' % max(y[-1] - exact_final)

# Test case 2: with res, adda, ia, ja & jac
m = method(res=res, adda_lsodis=adda, atol=atol, rtol=rtol, jac_lsodis=jac,
           ia=ia, ja=ja)
m.set_initial_condition(u0)
y, t = m.solve(time_points)
print 'Max error for test case 2 is %g' % max(y[-1] - exact_final)

# Test case 3: with res, adda & jac
m = method(res=res, adda_lsodis=adda, atol=atol, rtol=rtol, jac_lsodis=jac)
m.set_initial_condition(u0)
y, t = m.solve(time_points)
print 'Max error for test case 3 is %g' % max(y[-1] - exact_final)

# Test case 4:  With res & adda
m = method(res=res, adda_lsodis=adda, atol=atol, rtol=rtol,lrw=4000,liw=100)
m.set_initial_condition(u0)
y, t = m.solve(time_points)
print 'Max error for test case 4 is %g' % max(y[-1] - exact_final)


